import os
import pickle
from k_shortest_paths import single_source_ksp_yen

class Router(object):
    def __init__(self, map):
        self._map = map
        self._paths = {}
        
        data_file = "shortest_paths" + str(hash(map))+ ".dat"
        if os.path.exists(data_file):
            self.loadAllShortestPaths(data_file)
        else:
            self.findAllShortestPaths()
            self.saveAllShortestPaths(data_file)
        #print('all shortest paths', self._paths)
    
    def loadAllShortestPaths(self, data_file):
        """
        从文件中加载所有的最短路径。
        :param data_file: 最短路径所保存的文件。从该文件中加载最短路径。
        """
        with open(data_file,'rb') as f:
            self._paths = pickle.load(f)
        
    def saveAllShortestPaths(self, data_file):
        """
        将最短路径保存于文件中。
        :param data_file: 最短路径所保存的文件。
        """
        with open(data_file,'wb') as f:
            pickle.dump(self._paths, f)
        
    def findAllShortestPaths(self, max_k=3):
        """
        为所有的节点对查找max_k条最短路径。
        :param max_k: 为所有的节点对计算max_k条最短路径，默认值为3。
        """
        nodes = self._map.getAllNodes()
        for u in nodes:
            _, trails = single_source_ksp_yen(self._map.getGraph(), u, max_k=max_k)
            #print('trails', trails)
            self._paths.update(trails)
  
    def getAllShortestPaths(self):
        """
        获取所有最短路径。当最短路径没有被查找时，执行一次查找。若已执行过一次查找，或者已经从文件中加载，则直接返回结果。
        """
        if not self._paths:
            self.findAllShortestPaths()
        return self._paths

    def getKShortestPaths(self, src, dst):
        """
        获取从src到dst的（k条）最短路径。当最短路径没有被查找时，执行一次查找。若已执行过一次查找，或者已经从文件中加载，则直接返回结果。
        :param src: 起点。
        :param dst: 终点。
        """
        if not self._paths:
            self.findAllShortestPaths()
        return self._paths[(src, dst)]
        
    def getShortestPath(self, src, dst):
        """
        获取从src到dst的唯一的最短路径。当最短路径没有被查找时，执行一次查找。若已执行过一次查找，或者已经从文件中加载，则直接返回结果。
        :param src: 起点。
        :param dst: 终点。
        """
        if src == dst:
            return [src]
        
        if not self._paths:
            self.findAllShortestPaths()
            
        if (src, dst) in self._paths:
            return self._paths[(src, dst)][0]
        else:
            return []
        
    def getShortestPathLength(self, src, dst):
        """
        获取从src到dst的最短路径的长度。
        :param src: 起点。
        :param dst: 终点。
        """
        return len(self.getShortestPath(src, dst)) - 1
